import time
from typing import List, AnyStr

import simplejson as json
from pymongo import ASCENDING, IndexModel, MongoClient
from pymongo.errors import DuplicateKeyError

from .job import Job
from .base import BaseJobStore, JobLookupError, ConflictingIdError


class MongoDBJobStore(BaseJobStore):
    def __init__(self, db: str = 'python-jobs', collection: str = 'jobs', client: MongoClient = None, **connection_args):
        if not db:
            raise ValueError('db不能为空')
        if not collection:
            raise ValueError('collection不能为空')
        if client:
            self.client = client
        else:
            self.client = MongoClient(**connection_args)

        self.collection = self.client[db][collection]
        idx1 = IndexModel([('next_run_time', ASCENDING)])
        idx2 = IndexModel([('delete_time', ASCENDING)])
        idx3 = IndexModel([('_id', ASCENDING)])
        idx4 = IndexModel([('func_name', ASCENDING)])
        self.collection.create_indexes([idx1, idx2, idx3, idx4])

    def add_job(self, job: Job) -> None:
        '''
        :raises ConflictingIdError: 任务已经存在抛异常
        '''
        assert bool(job.job_id) is True
        assert isinstance(job.func_kwargs, dict)
        assert isinstance(job.func_args, (list, tuple))
        if not job.next_run_time:
            job.next_run_time = int(time.time()) + job.delay_seconds
        try:
            self.collection.insert({
                '_id': job.job_id,
                'func_name': job.func_name,
                'func_args': job.func_args,
                'func_kwargs': job.func_kwargs,
                'sched_times': job.sched_times,
                'max_sched_times': job.max_sched_times,
                'next_run_time': job.next_run_time,
                'delay_seconds': job.delay_seconds,
                'create_time': int(time.time()),
                'update_time': int(time.time()),
                'delete_time': 0,
            })
        except DuplicateKeyError:
            raise ConflictingIdError(job.job_id)

    def lookup_job(self, job_id: str) -> Job:
        document = self.collection.find_one({'_id': job_id})
        if not document:
            return None
        job = Job(**document)
        job.id = job.job_id = document['_id']
        return job

    def update_job(self, job: Job) -> None:
        '''
        :raises JobLookupError: 任务不存在抛异常
        '''
        assert isinstance(job, Job)
        changes = {
            'sched_times': job.sched_times,
            'update_time': int(time.time()),
            'next_run_time': job.next_run_time,
            'delete_time': job.delete_time
        }
        result = self.collection.update_one({'_id': job.job_id}, {'$set': changes})
        if result and result.modified_count == 0:
            raise JobLookupError(job.job_id)

    def remove_job(self, job: Job) -> None:
        assert isinstance(job, Job)
        result = self.collection.remove({'_id': job.job_id})
        if result and result['n'] == 0:
            raise JobLookupError(job.job_id)

    def get_due_jobs(self, now: int, count: int, local_job_ids: List[AnyStr] = []) -> List[Job]:
        '''
        :param int now: 时间戳
        :param int count: 获取任务数量
        '''
        conditions = {
            'next_run_time': {'$lte': now},
            'delete_time': 0,
            '$expr': {'$lt': ['$sched_times', '$max_sched_times']}
        }
        if len(local_job_ids) > 0:
            conditions['job_id'] = {'$nin': local_job_ids}
        documents = self.collection.find(conditions, sort=[('next_run_time', ASCENDING)]).limit(count)
        jobs = []
        for document in documents:
            job = Job(**document)
            job.id = job.job_id = document['_id']
            jobs.append(job)
        return jobs

    def get_next_run_time(self):
        '''
        :return int or None: 确保返回整数
        '''
        conditions = {
            'delete_time': 0,
            '$expr': {'$lt': ['$sched_times', '$max_sched_times']}
        }
        document = self.collection.find_one(
            conditions,
            projection=['next_run_time'],
            sort=[('next_run_time', ASCENDING)]
        )
        return int(document['next_run_time']) if document else None

    def lock_jobs(self, jobs: List[Job], lock_expire_time: int):
        job_ids = [job.job_id for job in jobs]
        self.collection.update_many({'_id': {'$in': job_ids}}, {'$set': {'next_run_time': lock_expire_time}})

    def remove_all_jobs(self):
        self.collection.remove({})
